import streamlit as st
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import io
from work1 import (
    sobel_edge_detection,
    given_kernel_filter_gray,
    compute_histogram_color
)

# 设置页面配置
st.set_page_config(
    page_title="图像处理应用",
    page_icon="🖼️",
    layout="wide"
)

# 自定义CSS样式
st.markdown("""
<style>
.main {
    padding: 2rem;
}
.stButton>button {
    width: 100%;
    background-color: #4CAF50;
    color: white;
    padding: 0.5rem;
    border-radius: 5px;
    border: none;
}
.stButton>button:hover {
    background-color: #45a049;
}
h1, h2, h3 {
    color: #2c3e50;
}
.stImage {
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
</style>
""", unsafe_allow_html=True)

# 标题
st.title("🎨 高级图像处理工具")
st.markdown("---")

# 文件上传
uploaded_file = st.file_uploader("选择一张图片进行处理", type=['png', 'jpg', 'jpeg'])


def plot_color_histogram(color_image):
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    colors = ['red', 'green', 'blue']
    histograms = compute_histogram_color(color_image)

    for idx, (ax, color, hist) in enumerate(zip(axes, colors, histograms)):
        ax.bar(range(256), hist, color=color, alpha=0.7)
        ax.set_title(f'{color.upper()} 通道直方图')
        ax.set_xlabel('像素值')
        ax.set_ylabel('频率')
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    return fig


def process_image(image):
    # 转换为numpy数组
    color_array = np.array(image)
    gray_array = np.array(image.convert('L'))

    # 处理图像
    sobel_result = sobel_edge_detection(gray_array)
    filtered_result = given_kernel_filter_gray(gray_array)

    return color_array, sobel_result, filtered_result


if uploaded_file is not None:
    # 读取图像
    image = Image.open(uploaded_file)

    # 创建三列布局
    col1, col2, col3 = st.columns(3)

    # 处理图像
    color_array, sobel_result, filtered_result = process_image(image)

    # 显示原图
    with col1:
        st.subheader("📸 原始图像")
        st.image(image, use_container_width=True)

    # 显示Sobel边缘检测结果
    with col2:
        st.subheader("🔍 Sobel边缘检测")
        st.image(sobel_result, use_container_width=True)

    # 显示给定卷积核滤波结果
    with col3:
        st.subheader("🎯 自定义卷积核滤波")
        st.image(filtered_result, use_container_width=True)

    # 显示颜色直方图
    st.subheader("📊 颜色直方图分析")
    hist_fig = plot_color_histogram(color_array)
    st.pyplot(hist_fig)

    # 添加下载按钮
    col1, col2, col3 = st.columns(3)


    # 转换图像为bytes
    def get_image_download_link(img, filename):
        buf = io.BytesIO()
        Image.fromarray(img).save(buf, format='PNG')
        return buf.getvalue()


    with col1:
        sobel_bytes = get_image_download_link(sobel_result, "sobel_result.png")
        st.download_button(
            label="📥 下载Sobel边缘检测结果",
            data=sobel_bytes,
            file_name="sobel_result.png",
            mime="image/png"
        )

    with col2:
        filtered_bytes = get_image_download_link(filtered_result, "filtered_result.png")
        st.download_button(
            label="📥 下载卷积核滤波结果",
            data=filtered_bytes,
            file_name="filtered_result.png",
            mime="image/png"
        )

else:
    # 显示欢迎信息
    st.info("👆 请上传一张图片开始处理")

    # 添加使用说明
    st.markdown("""
    ### 📝 使用说明
    1. 点击上方的"选择一张图片进行处理"按钮
    2. 选择您想要处理的图片文件（支持PNG、JPG、JPEG格式）
    3. 系统将自动处理并显示：
       - 原始图像
       - Sobel边缘检测结果
       - 自定义卷积核滤波结果
       - RGB通道的颜色直方图
    4. 您可以下载处理后的图像结果
    """)